// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <mutex>                           // std::lock_guard
#include <turbo/base/macros.h>
#include <turbo/log/logging.h>
#include <kthread/internal/errno.h>

namespace std {

#if defined(OS_POSIX)

    template<>
    class lock_guard<pthread_spinlock_t> {
    public:
        explicit lock_guard(pthread_spinlock_t &spin) : _pspin(&spin) {
#if !defined(NDEBUG)
            const int rc = pthread_spin_lock(_pspin);
            if (rc) {
                KLOG(FATAL) << "Fail to lock pthread_spinlock_t=" << _pspin << ", " << km_error(rc);
                _pspin = nullptr;
            }
#else
            pthread_spin_lock(_pspin);
#endif  // NDEBUG
        }

        ~lock_guard() {
#ifndef NDEBUG
            if (_pspin) {
                pthread_spin_unlock(_pspin);
            }
#else
            pthread_spin_unlock(_pspin);
#endif
        }

    private:
        TURBO_DISALLOW_COPY_AND_ASSIGN(lock_guard);

        pthread_spinlock_t *_pspin;
    };

    template<>
    class unique_lock<pthread_spinlock_t> {
        TURBO_DISALLOW_COPY_AND_ASSIGN(unique_lock);

    public:
        typedef pthread_spinlock_t mutex_type;

        unique_lock() : _mutex(nullptr), _owns_lock(false) {}

        explicit unique_lock(mutex_type &mutex)
                : _mutex(&mutex), _owns_lock(true) {
            pthread_spin_lock(_mutex);
        }

        ~unique_lock() {
            if (_owns_lock) {
                pthread_spin_unlock(_mutex);
            }
        }

        unique_lock(mutex_type &mutex, defer_lock_t)
                : _mutex(&mutex), _owns_lock(false) {}

        unique_lock(mutex_type &mutex, try_to_lock_t)
                : _mutex(&mutex), _owns_lock(pthread_spin_trylock(&mutex) == 0) {}

        unique_lock(mutex_type &mutex, adopt_lock_t)
                : _mutex(&mutex), _owns_lock(true) {}

        void lock() {
            if (_owns_lock) {
                KCHECK(false) << "Detected deadlock issue";
                return;
            }
#if !defined(NDEBUG)
            const int rc = pthread_spin_lock(_mutex);
            if (rc) {
                KLOG(FATAL) << "Fail to lock pthread_spinlock=" << _mutex << ", " << km_error(rc);
                return;
            }
            _owns_lock = true;
#else
            _owns_lock = true;
            pthread_spin_lock(_mutex);
#endif  // NDEBUG
        }

        bool try_lock() {
            if (_owns_lock) {
                KCHECK(false) << "Detected deadlock issue";
                return false;
            }
            _owns_lock = !pthread_spin_trylock(_mutex);
            return _owns_lock;
        }

        void unlock() {
            if (!_owns_lock) {
                KCHECK(false) << "Invalid operation";
                return;
            }
            pthread_spin_unlock(_mutex);
            _owns_lock = false;
        }

        void swap(unique_lock &rhs) {
            std::swap(_mutex, rhs._mutex);
            std::swap(_owns_lock, rhs._owns_lock);
        }

        mutex_type *release() {
            mutex_type *saved_mutex = _mutex;
            _mutex = nullptr;
            _owns_lock = false;
            return saved_mutex;
        }

        mutex_type *mutex() { return _mutex; }

        bool owns_lock() const { return _owns_lock; }

        operator bool() const { return owns_lock(); }

    private:
        mutex_type *_mutex;
        bool _owns_lock;
    };

#endif  // defined(OS_POSIX)

}  // namespace std

namespace kthread {

// Lock both lck1 and lck2 without the dead lock issue
    template<typename Mutex1, typename Mutex2>
    void double_lock(std::unique_lock<Mutex1> &lck1, std::unique_lock<Mutex2> &lck2) {
        DKCHECK(!lck1.owns_lock());
        DKCHECK(!lck2.owns_lock());
        volatile void *const ptr1 = lck1.mutex();
        volatile void *const ptr2 = lck2.mutex();
        DKCHECK_NE(ptr1, ptr2);
        if (ptr1 < ptr2) {
            lck1.lock();
            lck2.lock();
        } else {
            lck2.lock();
            lck1.lock();
        }
    }

}
